use clap::{Parser, Subcommand};
use std::str::FromStr;
use std::fs::File;
use std::io::{self, Write, BufRead, BufReader};
use std::path::PathBuf;
mod spline;
use spline::spline_struct::{Spline, Dot};
use nalgebra::{DMatrix, DVector};
mod solve_s;

#[derive(Parser, Debug)]
#[command(
    name = "spline_param_rs",
    author = "SplineParamRs Team",
    version,
    about = "A cubic spline parameterization tool",
    long_about = "Calculates cubic spline parameters for a set of points with optional boundary conditions"
)]
struct Cli {
    #[command(subcommand)]
    command: Option<Commands>,

    /// Points in format x1,y1 x2,y2 ... (if no subcommand is provided)
    #[arg(value_parser = parse_point)]
    points: Vec<Dot>,

    /// Left derivative value (一阶导数)
    #[arg(long = "left-prime", short = 'l', alias = "l1", visible_alias = "l1")]
    left_prime: Option<f64>,

    /// Left second derivative value (二阶导数)
    #[arg(long = "left-double-prime", short = 'L', alias = "l2", visible_alias = "l2")]
    left_double_prime: Option<f64>,

    /// Right derivative value (一阶导数)
    #[arg(long = "right-prime", short = 'r', alias = "r1", visible_alias = "r1")]
    right_prime: Option<f64>,

    /// Right second derivative value (二阶导数)
    #[arg(long = "right-double-prime", short = 'R', alias = "r2", visible_alias = "r2")]
    right_double_prime: Option<f64>,
}

#[derive(Subcommand, Debug)]
enum Commands {
    /// Calculate spline coefficients
    Solve {
        /// Points in format x1,y1 x2,y2 ...
        #[arg(required = true, value_parser = parse_point)]
        points: Vec<Dot>,

        /// Left derivative value (一阶导数)
        #[arg(long = "left-prime", short = 'l', alias = "l1", visible_alias = "l1")]
        left_prime: Option<f64>,

        /// Left second derivative value (二阶导数)
        #[arg(long = "left-double-prime", short = 'L', alias = "l2", visible_alias = "l2")]
        left_double_prime: Option<f64>,

        /// Right derivative value (一阶导数)
        #[arg(long = "right-prime", short = 'r', alias = "r1", visible_alias = "r1")]
        right_prime: Option<f64>,

        /// Right second derivative value (二阶导数)
        #[arg(long = "right-double-prime", short = 'R', alias = "r2", visible_alias = "r2")]
        right_double_prime: Option<f64>,

        /// Save coefficients to file
        #[arg(long, short)]
        output: Option<PathBuf>,
    },
    
    /// Evaluate spline at specific points
    Evaluate {
        /// Points to evaluate at (x values)
        #[arg(required = true)]
        x_values: Vec<f64>,
        
        /// File containing spline coefficients
        #[arg(long, short, required = true)]
        coefficients: PathBuf,
    },
    
    /// Save points to a file
    SavePoints {
        /// Points in format x1,y1 x2,y2 ...
        #[arg(required = true, value_parser = parse_point)]
        points: Vec<Dot>,
        
        /// Output file
        #[arg(long, short, required = true)]
        output: PathBuf,
    },
    
    /// Load points from a file
    LoadPoints {
        /// Input file containing points
        #[arg(required = true)]
        input: PathBuf,
        
        /// Left derivative value (一阶导数)
        #[arg(long = "left-prime", short = 'l', alias = "l1", visible_alias = "l1")]
        left_prime: Option<f64>,

        /// Left second derivative value (二阶导数)
        #[arg(long = "left-double-prime", short = 'L', alias = "l2", visible_alias = "l2")]
        left_double_prime: Option<f64>,

        /// Right derivative value (一阶导数)
        #[arg(long = "right-prime", short = 'r', alias = "r1", visible_alias = "r1")]
        right_prime: Option<f64>,

        /// Right second derivative value (二阶导数)
        #[arg(long = "right-double-prime", short = 'R', alias = "r2", visible_alias = "r2")]
        right_double_prime: Option<f64>,
        
        /// Save coefficients to file
        #[arg(long, short)]
        output: Option<PathBuf>,
    },
}

fn parse_point(s: &str) -> Result<Dot, String> {
    let parts: Vec<&str> = s.split(',').collect();
    if parts.len() != 2 {
        return Err("Points must be in format x,y".to_string());
    }
    let x = f64::from_str(parts[0]).map_err(|_| "Invalid x coordinate".to_string())?;
    let y = f64::from_str(parts[1]).map_err(|_| "Invalid y coordinate".to_string())?;
    Ok(Dot::new(x, y))
}

fn main() -> io::Result<()> {
    let cli = Cli::parse();
    
    match cli.command {
        Some(Commands::Solve { points, left_prime, left_double_prime, right_prime, right_double_prime, output }) => {
            solve_spline(points, left_prime, left_double_prime, right_prime, right_double_prime, output)?;
        },
        Some(Commands::Evaluate { x_values, coefficients }) => {
            evaluate_spline(x_values, coefficients)?;
        },
        Some(Commands::SavePoints { points, output }) => {
            save_points(points, output)?;
        },
        Some(Commands::LoadPoints { input, left_prime, left_double_prime, right_prime, right_double_prime, output }) => {
            load_and_solve_points(input, left_prime, left_double_prime, right_prime, right_double_prime, output)?;
        },
        None => {
            // Default behavior: solve with points from arguments
            if cli.points.is_empty() {
                println!("Error: No points provided. Use --help for usage information.");
                return Ok(());
            }
            solve_spline(cli.points, cli.left_prime, cli.left_double_prime, cli.right_prime, cli.right_double_prime, None)?;
        }
    }
    
    Ok(())
}

fn solve_spline(
    points: Vec<Dot>, 
    left_prime: Option<f64>, 
    left_double_prime: Option<f64>,
    right_prime: Option<f64>, 
    right_double_prime: Option<f64>,
    output: Option<PathBuf>
) -> io::Result<()> {
    let mut spline = Spline::new(points);
    
    if let Some(left) = left_prime {
        if let Err(len) = spline.give_left_prime(left) {
            println!("Error: Not enough points ({}). At least 3 points are required.", len);
            return Ok(());
        }
    }
    
    if let Some(left_double) = left_double_prime {
        if let Err(len) = spline.give_left_double_prime(left_double) {
            println!("Error: Not enough points ({}). At least 3 points are required.", len);
            return Ok(());
        }
    }
    
    if let Some(right) = right_prime {
        if let Err(len) = spline.give_right_prime(right) {
            println!("Error: Not enough points ({}). At least 3 points are required.", len);
            return Ok(());
        }
    }
    
    if let Some(right_double) = right_double_prime {
        if let Err(len) = spline.give_right_double_prime(right_double) {
            println!("Error: Not enough points ({}). At least 3 points are required.", len);
            return Ok(());
        }
    }
    
    match spline.get_matrix() {
        Ok((matrix, right)) => {
            let d_matrix = vec_to_dmatrix(matrix);
            let d_vector = vec_to_dvector(right);
            match solve_s::solve_linear_system(d_matrix, d_vector) {
                Ok(x) => {
                    print_coefficients(&spline.dots, &x);
                    
                    if let Some(path) = output {
                        save_coefficients(&spline.dots, &x, path)?;
                    }
                },
                Err(e) => {
                    println!("Error solving linear system: {:?}", e);
                }
            }
        },
        Err(e) => {
            println!("Error: Not enough boundary conditions ({}). Need exactly 2.", e);
        }
    }
    
    Ok(())
}

fn print_coefficients(dots: &[Dot], coefficients: &DVector<f64>) {
    println!("Spline Coefficients:");
    println!("-------------------");
    
    for i in 0..(dots.len() - 1) {
        let a = coefficients[4*i];
        let b = coefficients[4*i + 1];
        let c = coefficients[4*i + 2];
        let d = coefficients[4*i + 3];
        
        println!("Segment {} (x ∈ [{:.6}, {:.6}]):", i+1, dots[i].x, dots[i+1].x);
        println!("  f(x) = {:.6} + {:.6} x + {:.6} x² + {:.6} x³", 
            a, b, c, d);
    }
}

fn save_coefficients(dots: &[Dot], coefficients: &DVector<f64>, path: PathBuf) -> io::Result<()> {
    let mut file = File::create(path)?;
    
    writeln!(file, "# Spline Coefficients")?;
    writeln!(file, "# Format: segment_start segment_end a b c d")?;
    writeln!(file, "# f(x) = a + bx + cx^2 + dx^3")?;
    
    for i in 0..(dots.len() - 1) {
        let a = coefficients[4*i];
        let b = coefficients[4*i + 1];
        let c = coefficients[4*i + 2];
        let d = coefficients[4*i + 3];
        
        writeln!(file, "{} {} {} {} {} {}", 
            dots[i].x, dots[i+1].x, a, b, c, d)?;
    }
    
    println!("Coefficients saved successfully.");
    Ok(())
}

fn evaluate_spline(x_values: Vec<f64>, coefficients_path: PathBuf) -> io::Result<()> {
    let file = File::open(coefficients_path)?;
    let reader = BufReader::new(file);
    
    let mut segments = Vec::new();
    
    for line in reader.lines() {
        let line = line?;
        if line.starts_with('#') || line.trim().is_empty() {
            continue;
        }
        
        let parts: Vec<&str> = line.split_whitespace().collect();
        if parts.len() != 6 {
            println!("Error: Invalid coefficient file format");
            return Ok(());
        }
        
        let x0: f64 = parts[0].parse().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid x0"))?;
        let x1: f64 = parts[1].parse().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid x1"))?;
        let a: f64 = parts[2].parse().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid a"))?;
        let b: f64 = parts[3].parse().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid b"))?;
        let c: f64 = parts[4].parse().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid c"))?;
        let d: f64 = parts[5].parse().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid d"))?;
        
        segments.push((x0, x1, a, b, c, d));
    }
    
    println!("Evaluation Results:");
    println!("------------------");
    println!("    x        f(x)");
    
    for x in x_values {
        let mut found = false;
        
        for &(x0, x1, a, b, c, d) in &segments {
            if x >= x0 && x <= x1 {
                let dx = x - x0;
                let y = a + b*dx + c*dx*dx + d*dx*dx*dx;
                println!("{:8.4} {:8.4}", x, y);
                found = true;
                break;
            }
        }
        
        if !found {
            println!("{:8.4} Out of spline range", x);
        }
    }
    
    Ok(())
}

fn save_points(points: Vec<Dot>, path: PathBuf) -> io::Result<()> {
    let mut file = File::create(path)?;
    
    writeln!(file, "# Spline Points")?;
    writeln!(file, "# Format: x y")?;
    
    for point in points {
        writeln!(file, "{} {}", point.x, point.y)?;
    }
    
    println!("Points saved successfully.");
    Ok(())
}

fn load_and_solve_points(
    path: PathBuf, 
    left_prime: Option<f64>, 
    left_double_prime: Option<f64>,
    right_prime: Option<f64>, 
    right_double_prime: Option<f64>,
    output: Option<PathBuf>
) -> io::Result<()> {
    let file = File::open(path)?;
    let reader = BufReader::new(file);
    
    let mut points = Vec::new();
    
    for line in reader.lines() {
        let line = line?;
        if line.starts_with('#') || line.trim().is_empty() {
            continue;
        }
        
        let parts: Vec<&str> = line.split_whitespace().collect();
        if parts.len() != 2 {
            println!("Error: Invalid point file format");
            return Ok(());
        }
        
        let x: f64 = parts[0].parse().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid x"))?;
        let y: f64 = parts[1].parse().map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid y"))?;
        
        points.push(Dot::new(x, y));
    }
    
    println!("Loaded {} points", points.len());
    solve_spline(points, left_prime, left_double_prime, right_prime, right_double_prime, output)
}

fn vec_to_dmatrix(vec: Vec<Vec<f64>>) -> DMatrix<f64> {
    let mut matrix = DMatrix::zeros(vec.len(), vec[0].len());
    for i in 0..vec.len() {
        for j in 0..vec[0].len() {
            matrix[(i, j)] = vec[i][j];
        }
    }
    matrix
}

fn vec_to_dvector(vec: Vec<f64>) -> DVector<f64> {
    let mut vector = DVector::zeros(vec.len());
    for i in 0..vec.len() {
        vector[i] = vec[i];
    }
    vector
}